"""
This module provides classes and functions for an individual based modeling (ibm)
approach to the question of whether an MK-Wolbachia infected population is selected
for shifts in individual sex ratio (assumed to be under the control of females).
"""
__version__ = 0.04
__author__ = 'M. Flor'

import numpy as np
import numpy.random as npr
import matplotlib as mpl
import matplotlib.pyplot as plt
import copy, string, itertools, bisect, math, time, pprint, textwrap

def build_allele_names(n):
    return ["a{}".format(i) for i in range(n)]
    
def build_phenotype_map(alleles):
    return dict( zip(alleles,range(len(alleles))) )

def init_parameters(**kwargs):
    """Initialize parameters."""
    # first, fill parameter dictionary with default values:
    pars = {'female_fecundity': 10, \
            'mk_lethality': 1., \
            'mk_transmission_rate': 0.95, \
            'reallocation_efficiency': 0.2, \
            'mutation_rate': 1e-5, \
            'population_capacity': 100, \
            'base_fitness': 1.}
    # update default parameters with provided arguments:
    pars.update(**kwargs)
    # calculate the remaining parameters (these should not be provided explicitely!):
    pars['number_of_alleles'] = pars['female_fecundity'] + 1
    pars['alleles'] = build_allele_names(pars['number_of_alleles'])
    pars['phenotype_map'] = build_phenotype_map(pars['alleles'])
    return pars

# default module parameters:
_PARAMETERS = init_parameters()

def repr_parameters(pars):
    s = ''
    parnames = ['female_fecundity', 'mk_lethality', 'mk_transmission_rate', 'reallocation_efficiency', \
                'mutation_rate', 'population_capacity', 'base_fitness']
    for name in parnames:
        s += "{}: {}\n".format(name, pars[name])
    return s

def print_parameters(pars):
    print repr_parameters(pars)

def save_parameters(pars, filename='parameters.txt'):
    """Save parameters to file in a human readable format."""
    parfile = open(filename, 'w')
    parfile.write(repr_parameters(pars))
    parfile.close()
    
def load_parameters(filename):
    """Load parameters from file."""
    parfile = open(filename, 'r')
    pars = {}
    lines = parfile.readlines()
    for line in lines:
        p,v = line.split(':')
        pars[p] = eval(v)
    parfile.close()
    return initParameters(**pars)
     
def random_choice(seq):
    return seq[npr.randint(len(seq))]
    
def random_genotype(alleles):
    return [random_choice(alleles), random_choice(alleles)]

def shift_mutate(a, alleles):
    """Shift the mutated allele to a neighboring allele."""
    i = alleles.index(a)
    if npr.random() < 0.5:                # shift left
        if not i == 0:
            i -= 1
    else:                               # shift right
        if not i == len(alleles)-1:
            i += 1
    return alleles[i]

def jump_mutate(a, alleles):
    """Any allele different than the one to mutate is possible."""
    mutant_alleles = alleles[:]
    mutant_alleles.remove(a)
    return random_choice( mutant_alleles )
        
def mutate(a, alleles, mode='shift'):
    if mode=='jump':
        return jump_mutate(a, alleles)
    else:
        return shift_mutate(a, alleles)

def fitness_compensation(x, F, b, rmin=None):
    """
    Fitness compensation / resource reallocation based 
    
    x    - number of killed sons (integer value!)
    F    - total clutch size (female fecundity)
    rmin - resource threshold to viability
    b    - reallocation efficiency
    
    The fitness increase due to resource 
    Example:
    >>> from numpy import arange
    >>> from ibm_v4 import fitness_compensation
    >>> from scitools.aplotter import plot as aplot
    >>> F = 100
    >>> b = 0.5
    >>> rmin = 0.001    # default value
    >>> x = arange(0,F+1)
    >>> y = fitness_compensation(x, F, b, rmin)
    >>> aplot(x, y, dot='o', min_y=0, plot_slope=False)    # this plots fitness compensation per surviving offspring vs number of killed sons
       |                                                                            
       +10                                                                      o   
       |                                                                            
       |                                                                       o    
       |                                                                            
       |                                                                      o     
       |                                                                            
       |                                                                      o     
       |                                                                    oo      
       |                                                                    o       
       |                                                                  oo        
       |                                                                 o          
       |                                                              ooo           
       |                                                         ooooo              
       |                                                  oooooooo                  
       |                                 ooooooooooooooooo                          
       oooooooooooooooooooooooooooooooooo                                           
       |                                                                            
    ---+0-----------------------------------------------------------------------+---
       +0                                                                   +100    
    """
    if not rmin:
        rmin = 1./(10*F)            # default value: optimal resource level is 10 times the threshold
    a = (1-F*rmin)*(F-x)
    numerator   = a + b*x
    denominator = a + b*x*F*rmin
    return numerator / denominator

def fitness_compensation_old(x, b, t):
    """\
    x - sex ratio  (float value!)
    b - reallocation efficiency
    t - transmission rate of MK
    """
    return b*x*t/(1-x*t)

def treeindent(s, pos):
    lines = s.splitlines()
    prefix = '                |'+' '*(pos-5)
    inline = lines[0]+'\n'
    for l in lines[1:]:
        inline += "{}{}\n".format(prefix,l)
    return inline.rstrip()

def weighted_random_generator(weights, replacement=False):
    """Factory function to create a weighted random generator."""
    if replacement:
        return WeightedRandomGeneratorWithReplacement(weights)
    else:
        return WeightedRandomGeneratorWithoutReplacement(weights)

def random_individual(parameters=_PARAMETERS, mk=None):
    alleles = parameters['alleles']
    if npr.random() < 0.5:
        infection = None
    else:
        if mk == None:
            infection = Malekiller(parameters=parameters)
        else:
            infection = mk
    if npr.random() < 0.5:
        return Male(random_genotype(alleles), infection, parameters=parameters)
    else:
        return Female(random_genotype(alleles), infection, parameters=parameters)
        
def random_population(N, parameters=_PARAMETERS):
    males = []
    females = []
    mk = Malekiller(parameters=parameters)
    for i in range(N):
        individual = random_individual(parameters, mk)
        if individual.sex == 'male':
            males.append(individual)
        else:
            females.append(individual)
    return Population(males, females, parameters=parameters)

#def uniform_population(gnum=1, mk=Malekiller()):
#    """gnum - number of males/females per genotype"""
#    alleles = _ALLELES
#    nA = _NUMBER_OF_ALLELES
#    females = []
#    males = []
#    for i in range(nA):
#        for j in range(nA):
#            for k in range(gnum):
#                if npr.random() < 0.5:
#                    infection = None
#                else:
#                    infection = mk
#                male = Male([alleles[i],alleles[j]],infection)
#                if npr.random() < 0.5:
#                    infection = None
#                else:
#                    infection = mk
#                female = Female([alleles[i],alleles[j]],infection)
#                males.append(male)
#                females.append(female)
#    return Population(females,males,capacity=len(females+males))

class WeightedRandomGeneratorWithReplacement(object):
    def __init__(self, weights):
        self.totals = []
        running_total = 0

        for w in weights:
            running_total += w
            self.totals.append(running_total)

    def next(self):
        rnd = npr.random() * self.totals[-1]
        return bisect.bisect_right(self.totals, rnd)

    def __call__(self):
        return self.next()

class WeightedRandomGeneratorWithoutReplacement(WeightedRandomGeneratorWithReplacement):
    def __init__(self, weights):
        WeightedRandomGeneratorWithReplacement.__init__(self, weights)
        self.selected = []

    def next(self):
        if len(self.totals) == len(self.selected):
            raise ValueError, "generator is empty"
        rnd = npr.random() * self.totals[-1]
        i = bisect.bisect_right(self.totals, rnd)
        if i not in self.selected:
            self.selected.append(i)
            return i
        else:
            return self.next()

class Malekiller(object):
    def __init__(self, lethality=None, transmission_rate=None, parameters=_PARAMETERS):
        if lethality==None:
            self.lethality = parameters['mk_lethality']
        else:
            self.lethality = lethality
        if transmission_rate==None:
            self.transmission_rate = parameters['mk_transmission_rate']
        else:
            self.transmission_rate = transmission_rate
           
    def __str__(self):
        s = """\
            male killer
            -----------
               |-- lethality: {self.lethality}
               |-- transmission rate: {self.transmission_rate}
            """.format(self=self)
        return textwrap.dedent(s)
    
    def __repr__(self):
        return "Malekiller( lethality={self.lethality}, transmission_rate={self.transmission_rate} )".format(self=self)

class Individual(object):
    def __init__(self, genotype, infection=None, fitness=None, parameters=_PARAMETERS):
        self.genotype = genotype
        self.infection = self.mk = infection
        self.infected = self.infection is not None             # bool value so you can test for infection with 'if self.infected:...'
        if fitness == None:
            fitness = parameters['base_fitness']
        self.fitness = fitness
        self._alleles = parameters['alleles'] 
        self._phenomap = parameters['phenotype_map']
        self.phenotype = self.get_phenotypic_value()
        self._mutation_rate = parameters['mutation_rate']
        self._parameters = parameters
        
    def __str__(self):
        s = """\
            individual
            ----------
                |-- {infection}
                |-- phenotype: {self.phenotype}
                |-- fitness: {self.fitness}
                |-- genotype: {gtype}
            """.format( self=self, infection=self.infection2string(), gtype=self.genotype2string() )
        return textwrap.dedent(s)
        
    def __repr__(self):
        return "Individual( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, parameters={s._parameters!r})".format(s=self)        
    
    def infection2string(self):
        if self.infected:
            return "infected with {}".format( treeindent(str(self.infection),22) )
        else:
            return "uninfected"
    
    def genotype2string(self):
        return '{g[0]} || {g[1]}'.format(g=self.genotype)
        
    def get_phenotypic_value(self):
        """Take the average."""
        a1,a2 = self.genotype
        return ( self._phenomap[a1] + self._phenomap[a2] ) / 2.
    
    def update_fitness(self, factor):
        self.fitness *= factor
    
    def gamete(self):
        mutation = 0
        g = random_choice( self.genotype )
        # mutation:
        if npr.random() < self._mutation_rate:
            g = mutate(g, self._alleles, mode='shift')
            mutation = 1
        return g, mutation

class Female(Individual):
    def __init__(self, genotype, infection=None, fitness=None, parameters=_PARAMETERS):
        Individual.__init__(self, genotype=genotype, infection=infection, fitness=fitness, parameters=parameters)
        self.sex = 'female'
        self.sexratio = math.floor(self.phenotype)
        self.fecundity = parameters['female_fecundity']
        self._realloceff = parameters['reallocation_efficiency']
    
    def __repr__(self):
        return "Female( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, parameters={s._parameters!r})".format(s=self)
    
    def __str__(self):
        s = """\
            Female
            ------
                |-- {infection}
                |-- phenotype: {self.phenotype}
                |-- sex ratio: {self.sexratio}
                |-- fitness: {self.fitness}
                |-- genotype: {gtype}
            """.format( self=self, infection=self.infection2string(), gtype=self.genotype2string() )
        return textwrap.dedent(s)
    
    def produce_brood_with(self, male):
        mut_count = 0
        sons = []
        daughters = []
        pars = self._parameters
        for i in range(1,self.fecundity+1):
            # offspring genotype:
            egg, m1   = self.gamete()
            sperm, m2 = male.gamete()
            mut_count += m1+m2
            # offspring infection:
            infection = None
            if self.infected:
                if npr.random() < self.mk.transmission_rate:               # random transmission of infection
                    infection = self.mk
            # offspring sex:
            if i <= self.sexratio:
                sons.append( Male([egg,sperm], infection, parameters=pars) )
            else:
                daughters.append( Female([egg,sperm], infection, parameters=pars) )
        if self.infected:
            sons,daughters = self.male_killing(sons,daughters)
        male.number_of_matings += 1
        return sons, daughters, mut_count
    
    def male_killing(self, sons, daughters):
        n = len(sons)
        sons[:] = (s for s in sons if s.infected and not npr.random()<self.mk.lethality)
        body_count = n - len(sons)
        for offspring in sons+daughters:
            fc = fitness_compensation(body_count, self.fecundity, self._realloceff, rmin=None)
            offspring.update_fitness( fc )
        return sons,daughters
    
class Male(Individual):
    def __init__(self, genotype, infection=None, fitness=None, parameters=_PARAMETERS):
        Individual.__init__(self, genotype=genotype, infection=infection, fitness=fitness, parameters=parameters)
        self.sex = 'male'
        self.number_of_matings = 0
        
    def __str__(self):
        s = """\
            Male
            -----
                |-- {infection}
                |-- number of matings: {self.number_of_matings}
                |-- fitness: {self.fitness}
                |-- genotype: {gtype}
            """.format( self=self, infection=self.infection2string(), gtype=self.genotype2string() )
        return textwrap.dedent(s)
        
    def __repr__(self):
        return "Male( genotype={s.genotype}, infection={s.infection!r}, fitness={s.fitness}, parameters={s._parameters!r} )".format(s=self)

class Population(object):                             
    def __init__(self, males, females, generation=0, record=None, parameters=_PARAMETERS):
        self.males = males
        self.females = females
        self.size = len(males+females)        # distinguish between size and capacity
        self.generation = generation
        self.capacity = parameters['population_capacity']
        self._alleles = parameters['alleles']
        self._female_fecundity = parameters['female_fecundity']
        self._reallocation_efficiency = parameters['reallocation_efficiency']
        self._parameters = parameters
        
        if record == None:
            record_keys = ['generation','n','n_males','n_females','n_mk','n_mk_males','n_mk_females','n_mutations','average_sex_ratio','allele_count','genotype_count']
            self.record = {}
            for k in record_keys:
                self.record[k] = []
            self.update_record(mut_count=0)
        else:
            self.record = record
            
    def __str__(self):
        rec = {}                          # extract current state dictionary from record for easy access
        for k,l in self.record.items():
            rec[k] = l[-1]
        s = """\
            Population
            ----------
                |-- generation: {rec[generation]}
                |-- size: {rec[n]}  (capacity {s.capacity})
                |-- number of males: {rec[n_males]}  ({rec[n_mk_males]} infected)
                |-- number of females: {rec[n_females]}  ({rec[n_mk_females]} infected)
                |-- total infected: {rec[n_mk]}
                |-- average sex ratio: {rec[average_sex_ratio]}
            """.format( rec=rec, s=self )
        return textwrap.dedent(s)
    
    def allele_count(self, gcount=None):
        if gcount == None:
            gcount = self.genotype_count()
        return np.sum(gcount,0) + np.sum(gcount,1)
    
    def genotype_count(self):
        alleles = self._alleles
        N = len(alleles)
        gcount = np.zeros((N,N), dtype=int)
        for i in self.males+self.females:
            a1,a2 = i.genotype
            gcount[alleles.index(a1),alleles.index(a2)] += 1
        return gcount
    
    def average_sex_ratio(self):
        """Calculate average sex ratio (asr) among females."""
        value = 0.
        for f in self.females:
            value += f.sexratio
        return value/len(self.females)
    
    def mk_count(self):
        mk_males = mk_females = 0
        for m in self.males:
            if m.infected:
                mk_males += 1
        for f in self.females:
            if f.infected:
                mk_females += 1
        return mk_males, mk_females
    
    def tally(self):
        """Calculate average sex ratio, mk numbers, and genotype frequencies; but avoid iterating over population multiple times."""
        alleles = self._alleles
        nA = len(alleles)
        asr = 0.                                # average sex ratio among females
        mk_males = mk_females = 0               # numbers of infected males and females
        acount = [0 for a in alleles]           # allele count
        gcount = np.zeros((nA,nA), dtype=int)   # genotype count
        for m in self.males:
            if m.infected:
                mk_males += 1
            a1,a2 = m.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        for f in self.females:
            asr += f.sexratio
            if f.infected:
                mk_females += 1
            a1,a2 = f.genotype
            gcount[ alleles.index(a1), alleles.index(a2) ] += 1
        asr /= len(self.females)
        acount = self.allele_count(gcount)
        return asr, mk_males, mk_females, acount, gcount            
    
    def update_record(self, mut_count):
        try:
            last_stored = self.record['generation'][-1]
        except IndexError:                     # this Error is raised when the record is still empty
            last_stored = -1
        if self.generation != last_stored:     # only update record if the generation state has not been stored before
            self.record['generation'].append( self.generation )
            self.record['n'].append( self.size )
            self.record['n_males'].append( len(self.males) )
            self.record['n_females'].append( len(self.females) )
            self.record['n_mutations'].append( mut_count )
            sr, mkM, mkF, acount, gcount = self.tally()
            self.record['average_sex_ratio'].append( sr )
            self.record['n_mk_males'].append( mkM )
            self.record['n_mk_females'].append( mkF )
            self.record['n_mk'].append( mkM+mkF )
            self.record['allele_count'].append( acount )
            self.record['genotype_count'].append( gcount )
        
    def is_polymorphic(self):
        acount = self.record['allele_count'][-1]         # most current allele count
        maxcount = 2*self.capacity
        return maxcount not in acount
        
    def produce_offspring(self):
        mut_count = 0
        if self.males == []:
            raise ValueError, 'population extinct due to lack of males'
        if self.females == []:
            raise ValueError, 'population extinct due to lack of females'
            
        males = self.males[:]
        npr.shuffle(males)
        females = self.females[:]
        npr.shuffle(females)
        temp_male_pool = []            # after having mated, males go into this pool
        offspring = []                 # all offspring is collected here
        
        for f in females:
            try:
                m = males.pop()
            except IndexError:                  # if there are no unmated males left, 
                males = temp_male_pool[:]       # put all males from the temporary pool back
                temp_male_pool = []             # and empty the temporary pool again
                npr.shuffle(males)
                m = males.pop()
            sons,daughters,mut = f.produce_brood_with(m)
            mut_count += mut
            temp_male_pool.append(m)
            offspring.extend(sons+daughters)
        return offspring, mut_count
        
    def selection(self, offspring):
#        print "offspring size:", len(offspring)
        males = []
        females = []
        weights = []
        for o in offspring:
            weights.append(o.fitness)
        wrg = weighted_random_generator(weights, replacement=False)
        i = 0
        while i < self.capacity:
            try:
                individual = offspring[ wrg() ]
                if individual.sex == 'female':
                    females.append(individual)
                else:
                    males.append(individual)
                i += 1
            except IndexError:           # all offspring have made it into the next generation already (offspring number < population capacity)
                break
        return males, females
    
    def next(self):
        offspring, mut_count = self.produce_offspring()
        self.males, self.females = self.selection(offspring)
        self.size = len(self.males+self.females)
        self.generation += 1
        self.update_record(mut_count)
    
    def run(self, n, force=False, timing=False):
        """If force is False then n is maximal generation count but iteration stops if the population
        is no longer polymorphic."""
        if timing:
            g0 = self.generation
            t0 = time.time()
        if force:
            for i in range(n):
                self.next()
        else:
            g = 0
            while self.is_polymorphic() and g<n:
                self.next()
                g += 1
        if timing:
            print "Iterated {} generations in {:.2f} seconds".format(self.generation-g0, time.time()-t0)
            
    def visualize(self, plots=['asr','nM','mk','mut','acount','inds']):
        F = self._female_fecundity
        labels = dict(asr=dict(rec='average_sex_ratio',xlabel='generation',ylabel='average sex ratio',style='k-',ylims=[0,F]), \
                      nM=dict(rec='n_males',xlabel='generation',ylabel='number of males',style='g-',ylims=[0,self.capacity]), \
                      mk=dict(rec='n_mk',xlabel='generation',ylabel='MK infected',style='r-',ylims=[0,self.capacity]), \
                      mut=dict(rec='n_mutations',xlabel='generation',ylabel='number of mutations',style='m-'), \
                      acount=dict(rec='allele_count',xlabel='generation',ylabel='alleles'), \
                      inds=dict(rec='individuals',xlabel='allele 1',ylabel='allele 2') )
        if not isinstance(plots, list):
            plots = list(plots)
        if 'inds' in plots:
            plots.remove('inds')
        figs = []
        g = np.array(self.record['generation'])
        for p in plots:
            f = plt.figure()
            ax = f.add_subplot(111)
            figs.append(f)
            rec = labels[p]['rec']
            xl = labels[p]['xlabel']
            yl = labels[p]['ylabel']
            
            a = np.array(self.record[rec])
            if p == 'mut':
                a = np.cumsum(a)
                labels[p]['ylims']= [0,np.max(a)]
            if p == 'acount':
                a = np.transpose( np.array(a,dtype=float) )
                a /= (2*self.capacity)                         # normalize (so that values range between 0 and 1)
                ma = np.ma.array(a,mask=(a==0))            # mask zero values
                cmap = mpl.cm.jet                          # use jet color scheme
                cmap.set_bad('w')                          # use white for zero values
                im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
                cb = plt.colorbar(im)
                cb.set_label('allele frequency')
            elif p == 'inds':
                pass
            else:
                ylims = labels[p]['ylims']
                s = labels[p]['style']
                ax.plot(g,a,s)
                ax.set_ylim(*ylims)
            ax.set_xlabel(xl)
            ax.set_ylabel(yl)
        return plots,figs
        
    def genotype_show(self, g=-1):
        f = plt.figure()
        ax = f.add_subplot(111)
        a = np.array( self.record['genotype_count'][g], dtype=float )
        a /= self.capacity                                      # normalize
        ma = np.ma.array(a, mask=(a==0))
        cmap = mpl.cm.jet
        cmap.set_bad('w')
        im = ax.imshow(ma,cmap=cmap,vmin=0.,vmax=1.,origin='lower',interpolation='nearest',aspect='auto')
        cb = plt.colorbar(im)
        cb.set_label('frequency')
        ax.set_xlabel('chromosome 1')
        ax.set_ylabel('chromosome 2')
    
    def store_current_state(self, filename='population.dat'):
        pass
        
    
if __name__ == '__main__':
    npr.seed(21112016)
    
    pars = init_parameters(female_fecundity=10, mk_lethality=1., \
                           base_fitness=1., reallocation_efficiency=1., \
                           population_capacity=1000)
    print_parameters(pars)
    
    mk = Malekiller(parameters=pars)                      # use provided parameters (recommended usage)
#    mk2 = Malekiller()                                    # use module default parameters (good for debugging)
#    mk3 = Malekiller(lethality=0.5)                       # use default parameters but override lethality
#    mk4 = Malekiller(transmission_rate=0.25, parameters=pars)     # use provided parameters but override lethality
#    print mk, mk2, mk3, mk4
#    
#    f = Female(['a4','a7'], mk)
#    m = Male(['a5','a6'], mk)
#    print f
#    print m
#    sons,daughters,mut = f.produce_brood_with(m)
    
    pop = random_population(1000)
    print pop
    
    pop.run(100)
    print pop
    pop.visualize()
    plt.show()
#    pop = uniform_population()
##    females = [Female(['a1','a2'],'uninfected') for i in range(30)]
##    males = [Male(['a1','a2'],'uninfected') for i in range(30)]
##    pop = Population(females,males)
#    pop.run(10, force=True, timing=True)
#    
#    plots,figs = pop.visualize()    
#    plt.show()
